# Created by 黄景涛
# DATE 2024/5/7

from Cryptodome.Cipher import AES
from Cryptodome.Random import get_random_bytes
from Cryptodome.Util.Padding import pad, unpad


# 加密函数
def encrypt(plaintext: bytes, key: bytes, iv: bytes):
    # AES的密码器实例
    cipher = AES.new(key, AES.MODE_CBC, iv)
    # 不足一个AES块大小的，填充至一个块大小（16字节）
    padded_data = pad(plaintext, AES.block_size)
    # 调用密码器的加密方法，进行加密
    cipher_text = cipher.encrypt(padded_data)
    return cipher_text


# 解密函数
def decrypt(cipher_text: bytes, key: bytes, iv: bytes):
    """加密的逆过程"""
    # 创建一个相同效果的密码器
    cipher = AES.new(key, AES.MODE_CBC, iv)
    # 密文解密
    decrypted_text = cipher.decrypt(cipher_text)
    # 去除填充的字节
    s_text = unpad(decrypted_text, AES.block_size)
    return s_text


if __name__ == '__main__':
    # 加密密钥（16字节长度，对应128位）
    key = b'1234567890abcdef'
    # 初始向量（随机的）
    iv = get_random_bytes(16)
    print(f'密钥：{key}')
    print(f'初始向量：{iv}')
    # 明文
    password_text = 'Bk@123456'
    print(f'明文：{password_text}')
    text_encoding = password_text.encode('utf8')
    # plaintext = b'Hello, AES encryption!'

    # 加密明文
    cipher_text = encrypt(text_encoding, key, iv)
    print(f"密文：{cipher_text}")
    print(f"密文：{cipher_text.decode('latin1')}")

    # 解密密文
    decrypted_text = decrypt(cipher_text, key, iv)
    print(f"解密后：{decrypted_text.decode('utf-8')}")
